import json
import logging
from pathlib import Path

import numpy as np

from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
from lerobot.common.robot_devices.robots.configs import ManipulatorRobotConfig
from lerobot.common.robot_devices.robots.utils import get_arm_id
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
from lerobot.common.robot_devices.motors.feetech import TorqueMode


def ensure_safe_goal_position(
    goal_pos: np.ndarray, present_pos: np.ndarray, max_relative_target: float | list[float]
):
    # Cap relative action target magnitude for safety.
    diff = goal_pos - present_pos
    max_relative_target = np.array(max_relative_target)
    safe_diff = np.minimum(diff, max_relative_target)
    safe_diff = np.maximum(safe_diff, -max_relative_target)
    safe_goal_pos = present_pos + safe_diff

    if not np.allclose(goal_pos, safe_goal_pos):
        logging.warning(
            "Relative goal position magnitude had to be clamped to be safe.\n"
            f"  requested relative goal position target: {diff}\n"
            f"    clamped relative goal position target: {safe_diff}"
        )

    return safe_goal_pos


class ManipulatorRobot:
    def __init__(
        self,
        config: ManipulatorRobotConfig,
    ):
        self.config = config
        self.robot_type = self.config.type
        self.calibration_dir = Path(self.config.calibration_dir)
        # self.leader_arms = make_motors_buses_from_configs(self.config.leader_arms)
        self.follower_arms = make_motors_buses_from_configs(self.config.follower_arms)
        self.follower_arm_name = list(self.follower_arms.keys())[0]
        self.is_connected = False
        self.logs = {}

        self.joint_names = ['joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6']
        self.joint_velocity = [0.0] * len(self.joint_names)

    def get_motor_names(self, arm: dict[str, MotorsBus]) -> list:
        return [f"{arm}_{motor}" for arm, bus in arm.items() for motor in bus.motors]

    @property
    def motor_features(self) -> dict:
        action_names = self.get_motor_names(self.follower_arms)
        state_names = self.get_motor_names(self.follower_arms)
        return {
            "action": {
                "dtype": "float32",
                "shape": (len(action_names),),
                "names": action_names,
            },
            "observation.state": {
                "dtype": "float32",
                "shape": (len(state_names),),
                "names": state_names,
            },
        }

    @property
    def features(self):
        return {**self.motor_features}

    @property
    def available_arms(self):
        available_arms = []
        for name in self.follower_arms:
            arm_id = get_arm_id(name, "follower")
            available_arms.append(arm_id)
        # for name in self.leader_arms:
        #     arm_id = get_arm_id(name, "leader")
        #     available_arms.append(arm_id)
        return available_arms

    def connect(self):
        if self.is_connected:
            raise RobotDeviceAlreadyConnectedError(
                "ManipulatorRobot is already connected. Do not run `robot.connect()` twice."
            )

        if not self.follower_arms:
            raise ValueError(
                "ManipulatorRobot doesn't have any device to connect. See example of usage in docstring of the class."
            )

        # Connect the arms
        for name in self.follower_arms:
            print(f"Connecting {name} follower arm.")
            self.follower_arms[name].connect()
        # for name in self.leader_arms:
        #     print(f"Connecting {name} leader arm.")
        #     self.leader_arms[name].connect()


        # We assume that at connection time, arms are in a rest position, and torque can
        # be safely disabled to run calibration and/or set robot preset configurations.
        # for name in self.follower_arms:
        #     self.follower_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)
        # for name in self.leader_arms:
        #     self.leader_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)

        self.activate_calibration()

        if self.robot_type in ["so100", "so101", "moss", "lekiwi"]:
            self.set_so100_robot_preset()

        # Enable torque on all motors of the follower arms
        for name in self.follower_arms:
            print(f"Activating torque on {name} follower arm.")
            self.follower_arms[name].write("Torque_Enable", 1)

        # if self.config.gripper_open_degree is not None:
        #     if self.robot_type not in ["koch", "koch_bimanual"]:
        #         raise NotImplementedError(
        #             f"{self.robot_type} does not support position AND current control in the handle, which is require to set the gripper open."
        #         )
        #     # Set the leader arm in torque mode with the gripper motor set to an angle. This makes it possible
        #     # to squeeze the gripper and have it spring back to an open position on its own.
        #     for name in self.leader_arms:
        #         self.leader_arms[name].write("Torque_Enable", 1, "gripper")
        #         self.leader_arms[name].write("Goal_Position", self.config.gripper_open_degree, "gripper")

        # Check both arms can be read
        for name in self.follower_arms:
            self.follower_arms[name].read("Present_Position")
        # for name in self.leader_arms:
        #     self.leader_arms[name].read("Present_Position")

        self.is_connected = True

    def activate_calibration(self):
        """After calibration all motors function in human interpretable ranges.
        Rotations are expressed in degrees in nominal range of [-180, 180],
        and linear motions (like gripper of Aloha) in nominal range of [0, 100].
        """

        def load_or_run_calibration_(name, arm, arm_type):
            arm_id = get_arm_id(name, arm_type)
            arm_calib_path = self.calibration_dir / f"{arm_id}.json"

            if arm_calib_path.exists():
                with open(arm_calib_path) as f:
                    calibration = json.load(f)
            else:
                # TODO(rcadene): display a warning in __init__ if calibration file not available
                print(f"Missing calibration file '{arm_calib_path}'")

                if self.robot_type in ["koch", "koch_bimanual", "aloha"]:
                    from lerobot.common.robot_devices.robots.dynamixel_calibration import run_arm_calibration

                    calibration = run_arm_calibration(arm, self.robot_type, name, arm_type)

                elif self.robot_type in ["so100", "so101", "moss", "lekiwi"]:
                    from lerobot.common.robot_devices.robots.feetech_calibration import (
                        run_arm_manual_calibration,
                    )

                    calibration = run_arm_manual_calibration(arm, self.robot_type, name, arm_type)

                print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
                arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
                with open(arm_calib_path, "w") as f:
                    json.dump(calibration, f)

            return calibration

        for name, arm in self.follower_arms.items():
            calibration = load_or_run_calibration_(name, arm, "follower")
            arm.set_calibration(calibration)
        # for name, arm in self.leader_arms.items():
        #     calibration = load_or_run_calibration_(name, arm, "leader")
        #     arm.set_calibration(calibration)

    def set_so100_robot_preset(self):
        for name in self.follower_arms:
            # Mode=0 for Position Control
            self.follower_arms[name].write("Mode", 0)
            # Set P_Coefficient to lower value to avoid shakiness (Default is 32)
            self.follower_arms[name].write("P_Coefficient", 16)
            # Set I_Coefficient and D_Coefficient to default value 0 and 32
            self.follower_arms[name].write("I_Coefficient", 0)
            self.follower_arms[name].write("D_Coefficient", 32)
            # Close the write lock so that Maximum_Acceleration gets written to EPROM address,
            # which is mandatory for Maximum_Acceleration to take effect after rebooting.
            self.follower_arms[name].write("Lock", 0)
            # Set Maximum_Acceleration to 254 to speedup acceleration and deceleration of
            # the motors. Note: this configuration is not in the official STS3215 Memory Table
            self.follower_arms[name].write("Maximum_Acceleration", 254)
            self.follower_arms[name].write("Acceleration", 254)
    
    def get_positions(self, degree=False):
        positions = {}
        for name in self.follower_arms:
            positions[name] = self.follower_arms[name].read("Present_Position")
        positions = positions[name]
        if not degree:
            positions = np.deg2rad(positions)
        return positions
    
    def set_positions(self, positions, degree=False):
        goal_pos = np.array(positions)
        gripper_pos = goal_pos[-1]              # 夹爪是百分比，特殊处理
        if not degree:
            goal_pos = np.rad2deg(goal_pos)
        if self.config.max_relative_target is not None:
            present_pos = self.get_positions(degree=True)
            goal_pos = ensure_safe_goal_position(goal_pos, present_pos, self.config.max_relative_target)
        goal_pos[-1] = gripper_pos
        goal_pos = goal_pos.astype(np.float32)
        self.follower_arms[self.follower_arm_name].write("Goal_Position", goal_pos)
    
    def enable_motor(self):
        for name in self.follower_arms:
            print(f"Activating torque on {name} follower arm.")
            self.follower_arms[name].write("Torque_Enable", 1)

    def disable_motor(self):
        for name in self.follower_arms:
            print(f"Activating torque off {name} follower arm.")
            self.follower_arms[name].write("Torque_Enable", 0)

    def print_logs(self):
        pass
        # TODO(aliberts): move robot-specific logs logic here

    def disconnect(self):
        if not self.is_connected:
            raise RobotDeviceNotConnectedError(
                "ManipulatorRobot is not connected. You need to run `robot.connect()` before disconnecting."
            )

        for name in self.follower_arms:
            self.follower_arms[name].disconnect()

        # for name in self.leader_arms:
        #     self.leader_arms[name].disconnect()

        # for name in self.cameras:
        #     self.cameras[name].disconnect()

        self.is_connected = False

    def __del__(self):
        if getattr(self, "is_connected", False):
            self.disconnect()
